Repository navigation
Conversation
klin2024
marked this pull request as ready for review
November 9, 2026 15:01
gate_msa / c_gate_msa are float32, so `hidden_states + attn_output * gate_msa` is already computed in float32 through type promotion; the explicit .float() calls only made full-size float32 copies of hidden_states and ff_output. Same change in the Wan-derived blocks (Chronoedit and Wan-Animate via make fix-copies, Helios, AnyFlow, AnyFlow-FAR, SkyReels-V2, Wan-VACE control block). Bit-identical output and gradients. Signed-off-by: Lin, KuanTing <Kuanting.Lin@amd.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #15224
The Wan block computes its gated residual adds as
(hidden_states.float() + attn_output * gate_msa)and(hidden_states.float() + ff_output.float() * c_gate_msa).gate_msa/c_gate_msaare already float32 (scale_shift_table + temb.float()), so type promotion does these adds in float32 anyway; the.float()calls only make full-size float32 copies in every block. This PR drops them. Output and gradients are bit-identical.Same change in the Wan-derived blocks: Chronoedit and Wan-Animate (
# Copied from, viamake fix-copies), Helios, AnyFlow, AnyFlow-FAR, SkyReels-V2 and the Wan-VACE control block.Transformer forward, eager, bf16, before -> after (ms):
RTX 5090: torch 2.14.1+cu130. R9700: torch 2.15.0a0+rocm10.1. diffusers main 8c54b38. Real transformer configs with random weights (timing does not depend on weight values), median of 6 forwards after warmup, 3 interleaved rounds per side; outputs bit-identical in every round on both GPUs. Gradients checked on the residual expression itself (with
hidden_statesalso feeding a second branch, as in the block): bit-identical for bf16 / fp16 on CUDA, ROCm and CPU. (The model-level backward is not deterministic run to run on main already, so it cannot be compared bit for bit.)Tests
Transformer tests of the touched models (Wan, Wan-VACE, Wan-Animate, Wan-Animate-2, Chronoedit, Helios, AnyFlow, AnyFlow-FAR, SkyReels-V2 and related, 12 files; GGUF / single-file / BFL tests deselected since they download full checkpoints), RTX 5090: 5 failed, 504 passed, 216 skipped. Main: the same counts and the same 5
test_compile_works_with_aotfailures.Self-review notes (
self-reviewskill)READY, no blocking issues. Left as is:
self.modulation + tembwithout.float(), which is float32 only whenmodulationstays in float32 (_keep_in_fp32_modulesviafrom_pretrained); after a plain.to(torch.bfloat16)dropping the cast would change the result.Before submitting
self-reviewskill on the diff?cc @sayakpaul @dg845